// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009-2010 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_GENERAL_MATRIX_MATRIX_TRIANGULAR_H
#define EIGEN_GENERAL_MATRIX_MATRIX_TRIANGULAR_H

namespace Eigen {

template <typename Scalar, typename Index, int StorageOrder, int UpLo, bool ConjLhs, bool ConjRhs> struct selfadjoint_rank1_update;

namespace internal {

    /**********************************************************************
* This file implements a general A * B product while
* evaluating only one triangular part of the product.
* This is a more general version of self adjoint product (C += A A^T)
* as the level 3 SYRK Blas routine.
**********************************************************************/

    // forward declarations (defined at the end of this file)
    template <typename LhsScalar, typename RhsScalar, typename Index, int mr, int nr, bool ConjLhs, bool ConjRhs, int ResInnerStride, int UpLo>
    struct tribb_kernel;

    /* Optimized matrix-matrix product evaluating only one triangular half */
    template <typename Index,
              typename LhsScalar,
              int LhsStorageOrder,
              bool ConjugateLhs,
              typename RhsScalar,
              int RhsStorageOrder,
              bool ConjugateRhs,
              int ResStorageOrder,
              int ResInnerStride,
              int UpLo,
              int Version = Specialized>
    struct general_matrix_matrix_triangular_product;

    // as usual if the result is row major => we transpose the product
    template <typename Index,
              typename LhsScalar,
              int LhsStorageOrder,
              bool ConjugateLhs,
              typename RhsScalar,
              int RhsStorageOrder,
              bool ConjugateRhs,
              int ResInnerStride,
              int UpLo,
              int Version>
    struct general_matrix_matrix_triangular_product<Index,
                                                    LhsScalar,
                                                    LhsStorageOrder,
                                                    ConjugateLhs,
                                                    RhsScalar,
                                                    RhsStorageOrder,
                                                    ConjugateRhs,
                                                    RowMajor,
                                                    ResInnerStride,
                                                    UpLo,
                                                    Version>
    {
        typedef typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType ResScalar;
        static EIGEN_STRONG_INLINE void run(Index size,
                                            Index depth,
                                            const LhsScalar* lhs,
                                            Index lhsStride,
                                            const RhsScalar* rhs,
                                            Index rhsStride,
                                            ResScalar* res,
                                            Index resIncr,
                                            Index resStride,
                                            const ResScalar& alpha,
                                            level3_blocking<RhsScalar, LhsScalar>& blocking)
        {
            general_matrix_matrix_triangular_product<
                Index,
                RhsScalar,
                RhsStorageOrder == RowMajor ? ColMajor : RowMajor,
                ConjugateRhs,
                LhsScalar,
                LhsStorageOrder == RowMajor ? ColMajor : RowMajor,
                ConjugateLhs,
                ColMajor,
                ResInnerStride,
                UpLo == Lower ? Upper : Lower>::run(size, depth, rhs, rhsStride, lhs, lhsStride, res, resIncr, resStride, alpha, blocking);
        }
    };

    template <typename Index,
              typename LhsScalar,
              int LhsStorageOrder,
              bool ConjugateLhs,
              typename RhsScalar,
              int RhsStorageOrder,
              bool ConjugateRhs,
              int ResInnerStride,
              int UpLo,
              int Version>
    struct general_matrix_matrix_triangular_product<Index,
                                                    LhsScalar,
                                                    LhsStorageOrder,
                                                    ConjugateLhs,
                                                    RhsScalar,
                                                    RhsStorageOrder,
                                                    ConjugateRhs,
                                                    ColMajor,
                                                    ResInnerStride,
                                                    UpLo,
                                                    Version>
    {
        typedef typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType ResScalar;
        static EIGEN_STRONG_INLINE void run(Index size,
                                            Index depth,
                                            const LhsScalar* _lhs,
                                            Index lhsStride,
                                            const RhsScalar* _rhs,
                                            Index rhsStride,
                                            ResScalar* _res,
                                            Index resIncr,
                                            Index resStride,
                                            const ResScalar& alpha,
                                            level3_blocking<LhsScalar, RhsScalar>& blocking)
        {
            typedef gebp_traits<LhsScalar, RhsScalar> Traits;

            typedef const_blas_data_mapper<LhsScalar, Index, LhsStorageOrder> LhsMapper;
            typedef const_blas_data_mapper<RhsScalar, Index, RhsStorageOrder> RhsMapper;
            typedef blas_data_mapper<typename Traits::ResScalar, Index, ColMajor, Unaligned, ResInnerStride> ResMapper;
            LhsMapper lhs(_lhs, lhsStride);
            RhsMapper rhs(_rhs, rhsStride);
            ResMapper res(_res, resStride, resIncr);

            Index kc = blocking.kc();
            Index mc = (std::min)(size, blocking.mc());

            // !!! mc must be a multiple of nr:
            if (mc > Traits::nr)
                mc = (mc / Traits::nr) * Traits::nr;

            std::size_t sizeA = kc * mc;
            std::size_t sizeB = kc * size;

            ei_declare_aligned_stack_constructed_variable(LhsScalar, blockA, sizeA, blocking.blockA());
            ei_declare_aligned_stack_constructed_variable(RhsScalar, blockB, sizeB, blocking.blockB());

            gemm_pack_lhs<LhsScalar, Index, LhsMapper, Traits::mr, Traits::LhsProgress, typename Traits::LhsPacket4Packing, LhsStorageOrder> pack_lhs;
            gemm_pack_rhs<RhsScalar, Index, RhsMapper, Traits::nr, RhsStorageOrder> pack_rhs;
            gebp_kernel<LhsScalar, RhsScalar, Index, ResMapper, Traits::mr, Traits::nr, ConjugateLhs, ConjugateRhs> gebp;
            tribb_kernel<LhsScalar, RhsScalar, Index, Traits::mr, Traits::nr, ConjugateLhs, ConjugateRhs, ResInnerStride, UpLo> sybb;

            for (Index k2 = 0; k2 < depth; k2 += kc)
            {
                const Index actual_kc = (std::min)(k2 + kc, depth) - k2;

                // note that the actual rhs is the transpose/adjoint of mat
                pack_rhs(blockB, rhs.getSubMapper(k2, 0), actual_kc, size);

                for (Index i2 = 0; i2 < size; i2 += mc)
                {
                    const Index actual_mc = (std::min)(i2 + mc, size) - i2;

                    pack_lhs(blockA, lhs.getSubMapper(i2, k2), actual_kc, actual_mc);

                    // the selected actual_mc * size panel of res is split into three different part:
                    //  1 - before the diagonal => processed with gebp or skipped
                    //  2 - the actual_mc x actual_mc symmetric block => processed with a special kernel
                    //  3 - after the diagonal => processed with gebp or skipped
                    if (UpLo == Lower)
                        gebp(res.getSubMapper(i2, 0), blockA, blockB, actual_mc, actual_kc, (std::min)(size, i2), alpha, -1, -1, 0, 0);

                    sybb(_res + resStride * i2 + resIncr * i2, resIncr, resStride, blockA, blockB + actual_kc * i2, actual_mc, actual_kc, alpha);

                    if (UpLo == Upper)
                    {
                        Index j2 = i2 + actual_mc;
                        gebp(res.getSubMapper(i2, j2),
                             blockA,
                             blockB + actual_kc * j2,
                             actual_mc,
                             actual_kc,
                             (std::max)(Index(0), size - j2),
                             alpha,
                             -1,
                             -1,
                             0,
                             0);
                    }
                }
            }
        }
    };

    // Optimized packed Block * packed Block product kernel evaluating only one given triangular part
    // This kernel is built on top of the gebp kernel:
    // - the current destination block is processed per panel of actual_mc x BlockSize
    //   where BlockSize is set to the minimal value allowing gebp to be as fast as possible
    // - then, as usual, each panel is split into three parts along the diagonal,
    //   the sub blocks above and below the diagonal are processed as usual,
    //   while the triangular block overlapping the diagonal is evaluated into a
    //   small temporary buffer which is then accumulated into the result using a
    //   triangular traversal.
    template <typename LhsScalar, typename RhsScalar, typename Index, int mr, int nr, bool ConjLhs, bool ConjRhs, int ResInnerStride, int UpLo>
    struct tribb_kernel
    {
        typedef gebp_traits<LhsScalar, RhsScalar, ConjLhs, ConjRhs> Traits;
        typedef typename Traits::ResScalar ResScalar;

        enum
        {
            BlockSize = meta_least_common_multiple<EIGEN_PLAIN_ENUM_MAX(mr, nr), EIGEN_PLAIN_ENUM_MIN(mr, nr)>::ret
        };
        void operator()(ResScalar* _res,
                        Index resIncr,
                        Index resStride,
                        const LhsScalar* blockA,
                        const RhsScalar* blockB,
                        Index size,
                        Index depth,
                        const ResScalar& alpha)
        {
            typedef blas_data_mapper<ResScalar, Index, ColMajor, Unaligned, ResInnerStride> ResMapper;
            typedef blas_data_mapper<ResScalar, Index, ColMajor, Unaligned> BufferMapper;
            ResMapper res(_res, resStride, resIncr);
            gebp_kernel<LhsScalar, RhsScalar, Index, ResMapper, mr, nr, ConjLhs, ConjRhs> gebp_kernel1;
            gebp_kernel<LhsScalar, RhsScalar, Index, BufferMapper, mr, nr, ConjLhs, ConjRhs> gebp_kernel2;

            Matrix<ResScalar, BlockSize, BlockSize, ColMajor> buffer((internal::constructor_without_unaligned_array_assert()));

            // let's process the block per panel of actual_mc x BlockSize,
            // again, each is split into three parts, etc.
            for (Index j = 0; j < size; j += BlockSize)
            {
                Index actualBlockSize = std::min<Index>(BlockSize, size - j);
                const RhsScalar* actual_b = blockB + j * depth;

                if (UpLo == Upper)
                    gebp_kernel1(res.getSubMapper(0, j), blockA, actual_b, j, depth, actualBlockSize, alpha, -1, -1, 0, 0);

                // selfadjoint micro block
                {
                    Index i = j;
                    buffer.setZero();
                    // 1 - apply the kernel on the temporary buffer
                    gebp_kernel2(
                        BufferMapper(buffer.data(), BlockSize), blockA + depth * i, actual_b, actualBlockSize, depth, actualBlockSize, alpha, -1, -1, 0, 0);

                    // 2 - triangular accumulation
                    for (Index j1 = 0; j1 < actualBlockSize; ++j1)
                    {
                        typename ResMapper::LinearMapper r = res.getLinearMapper(i, j + j1);
                        for (Index i1 = UpLo == Lower ? j1 : 0; UpLo == Lower ? i1 < actualBlockSize : i1 <= j1; ++i1) r(i1) += buffer(i1, j1);
                    }
                }

                if (UpLo == Lower)
                {
                    Index i = j + actualBlockSize;
                    gebp_kernel1(res.getSubMapper(i, j), blockA + depth * i, actual_b, size - i, depth, actualBlockSize, alpha, -1, -1, 0, 0);
                }
            }
        }
    };

}  // end namespace internal

// high level API

template <typename MatrixType, typename ProductType, int UpLo, bool IsOuterProduct> struct general_product_to_triangular_selector;

template <typename MatrixType, typename ProductType, int UpLo> struct general_product_to_triangular_selector<MatrixType, ProductType, UpLo, true>
{
    static void run(MatrixType& mat, const ProductType& prod, const typename MatrixType::Scalar& alpha, bool beta)
    {
        typedef typename MatrixType::Scalar Scalar;

        typedef typename internal::remove_all<typename ProductType::LhsNested>::type Lhs;
        typedef internal::blas_traits<Lhs> LhsBlasTraits;
        typedef typename LhsBlasTraits::DirectLinearAccessType ActualLhs;
        typedef typename internal::remove_all<ActualLhs>::type _ActualLhs;
        typename internal::add_const_on_value_type<ActualLhs>::type actualLhs = LhsBlasTraits::extract(prod.lhs());

        typedef typename internal::remove_all<typename ProductType::RhsNested>::type Rhs;
        typedef internal::blas_traits<Rhs> RhsBlasTraits;
        typedef typename RhsBlasTraits::DirectLinearAccessType ActualRhs;
        typedef typename internal::remove_all<ActualRhs>::type _ActualRhs;
        typename internal::add_const_on_value_type<ActualRhs>::type actualRhs = RhsBlasTraits::extract(prod.rhs());

        Scalar actualAlpha = alpha * LhsBlasTraits::extractScalarFactor(prod.lhs().derived()) * RhsBlasTraits::extractScalarFactor(prod.rhs().derived());

        if (!beta)
            mat.template triangularView<UpLo>().setZero();

        enum
        {
            StorageOrder = (internal::traits<MatrixType>::Flags & RowMajorBit) ? RowMajor : ColMajor,
            UseLhsDirectly = _ActualLhs::InnerStrideAtCompileTime == 1,
            UseRhsDirectly = _ActualRhs::InnerStrideAtCompileTime == 1
        };

        internal::gemv_static_vector_if<Scalar, Lhs::SizeAtCompileTime, Lhs::MaxSizeAtCompileTime, !UseLhsDirectly> static_lhs;
        ei_declare_aligned_stack_constructed_variable(
            Scalar, actualLhsPtr, actualLhs.size(), (UseLhsDirectly ? const_cast<Scalar*>(actualLhs.data()) : static_lhs.data()));
        if (!UseLhsDirectly)
            Map<typename _ActualLhs::PlainObject>(actualLhsPtr, actualLhs.size()) = actualLhs;

        internal::gemv_static_vector_if<Scalar, Rhs::SizeAtCompileTime, Rhs::MaxSizeAtCompileTime, !UseRhsDirectly> static_rhs;
        ei_declare_aligned_stack_constructed_variable(
            Scalar, actualRhsPtr, actualRhs.size(), (UseRhsDirectly ? const_cast<Scalar*>(actualRhs.data()) : static_rhs.data()));
        if (!UseRhsDirectly)
            Map<typename _ActualRhs::PlainObject>(actualRhsPtr, actualRhs.size()) = actualRhs;

        selfadjoint_rank1_update<Scalar,
                                 Index,
                                 StorageOrder,
                                 UpLo,
                                 LhsBlasTraits::NeedToConjugate && NumTraits<Scalar>::IsComplex,
                                 RhsBlasTraits::NeedToConjugate && NumTraits<Scalar>::IsComplex>::run(actualLhs.size(),
                                                                                                      mat.data(),
                                                                                                      mat.outerStride(),
                                                                                                      actualLhsPtr,
                                                                                                      actualRhsPtr,
                                                                                                      actualAlpha);
    }
};

template <typename MatrixType, typename ProductType, int UpLo> struct general_product_to_triangular_selector<MatrixType, ProductType, UpLo, false>
{
    static void run(MatrixType& mat, const ProductType& prod, const typename MatrixType::Scalar& alpha, bool beta)
    {
        typedef typename internal::remove_all<typename ProductType::LhsNested>::type Lhs;
        typedef internal::blas_traits<Lhs> LhsBlasTraits;
        typedef typename LhsBlasTraits::DirectLinearAccessType ActualLhs;
        typedef typename internal::remove_all<ActualLhs>::type _ActualLhs;
        typename internal::add_const_on_value_type<ActualLhs>::type actualLhs = LhsBlasTraits::extract(prod.lhs());

        typedef typename internal::remove_all<typename ProductType::RhsNested>::type Rhs;
        typedef internal::blas_traits<Rhs> RhsBlasTraits;
        typedef typename RhsBlasTraits::DirectLinearAccessType ActualRhs;
        typedef typename internal::remove_all<ActualRhs>::type _ActualRhs;
        typename internal::add_const_on_value_type<ActualRhs>::type actualRhs = RhsBlasTraits::extract(prod.rhs());

        typename ProductType::Scalar actualAlpha =
            alpha * LhsBlasTraits::extractScalarFactor(prod.lhs().derived()) * RhsBlasTraits::extractScalarFactor(prod.rhs().derived());

        if (!beta)
            mat.template triangularView<UpLo>().setZero();

        enum
        {
            IsRowMajor = (internal::traits<MatrixType>::Flags & RowMajorBit) ? 1 : 0,
            LhsIsRowMajor = _ActualLhs::Flags & RowMajorBit ? 1 : 0,
            RhsIsRowMajor = _ActualRhs::Flags & RowMajorBit ? 1 : 0,
            SkipDiag = (UpLo & (UnitDiag | ZeroDiag)) != 0
        };

        Index size = mat.cols();
        if (SkipDiag)
            size--;
        Index depth = actualLhs.cols();

        typedef internal::gemm_blocking_space<IsRowMajor ? RowMajor : ColMajor,
                                              typename Lhs::Scalar,
                                              typename Rhs::Scalar,
                                              MatrixType::MaxColsAtCompileTime,
                                              MatrixType::MaxColsAtCompileTime,
                                              _ActualRhs::MaxColsAtCompileTime>
            BlockingType;

        BlockingType blocking(size, size, depth, 1, false);

        internal::general_matrix_matrix_triangular_product<Index,
                                                           typename Lhs::Scalar,
                                                           LhsIsRowMajor ? RowMajor : ColMajor,
                                                           LhsBlasTraits::NeedToConjugate,
                                                           typename Rhs::Scalar,
                                                           RhsIsRowMajor ? RowMajor : ColMajor,
                                                           RhsBlasTraits::NeedToConjugate,
                                                           IsRowMajor ? RowMajor : ColMajor,
                                                           MatrixType::InnerStrideAtCompileTime,
                                                           UpLo&(Lower | Upper)>::run(size,
                                                                                      depth,
                                                                                      &actualLhs.coeffRef(SkipDiag && (UpLo & Lower) == Lower ? 1 : 0, 0),
                                                                                      actualLhs.outerStride(),
                                                                                      &actualRhs.coeffRef(0, SkipDiag && (UpLo & Upper) == Upper ? 1 : 0),
                                                                                      actualRhs.outerStride(),
                                                                                      mat.data() + (SkipDiag ? (bool(IsRowMajor) != ((UpLo & Lower) == Lower) ?
                                                                                                                    mat.innerStride() :
                                                                                                                    mat.outerStride()) :
                                                                                                               0),
                                                                                      mat.innerStride(),
                                                                                      mat.outerStride(),
                                                                                      actualAlpha,
                                                                                      blocking);
    }
};

template <typename MatrixType, unsigned int UpLo>
template <typename ProductType>
EIGEN_DEVICE_FUNC TriangularView<MatrixType, UpLo>&
TriangularViewImpl<MatrixType, UpLo, Dense>::_assignProduct(const ProductType& prod, const Scalar& alpha, bool beta)
{
    EIGEN_STATIC_ASSERT((UpLo & UnitDiag) == 0, WRITING_TO_TRIANGULAR_PART_WITH_UNIT_DIAGONAL_IS_NOT_SUPPORTED);
    eigen_assert(derived().nestedExpression().rows() == prod.rows() && derived().cols() == prod.cols());

    general_product_to_triangular_selector<MatrixType, ProductType, UpLo, internal::traits<ProductType>::InnerSize == 1>::run(
        derived().nestedExpression().const_cast_derived(), prod, alpha, beta);

    return derived();
}

}  // end namespace Eigen

#endif  // EIGEN_GENERAL_MATRIX_MATRIX_TRIANGULAR_H
